/*
 * Speciation.cpp
 *
 *  Created on: 13 Jun 2011
 *      Author: allan
 */

#include "Speciation.h"

// C++ includes
#include <iostream>
#include <iomanip>

// GeoReact includes
#include "MultirootNewtonRaphson.h"
#include "NumericalUtilities.h"
#include "Utils.h"

Speciation::Speciation()
{
}

Speciation::~Speciation()
{
}

void Speciation::SetSystem(System& system)
{
	this->system = &system;
	
	// the number of primary species
	const unsigned numPrimary = system.GetNumOfPrimarySpecies();
	
	// the number of equilibrium species
	const unsigned numEquilibrium = system.GetNumOfEquilibriumSpecies();
	
	zj.resize(numPrimary);
	
	for(unsigned iPrimary = 0; iPrimary < numPrimary; ++iPrimary)
		zj[iPrimary] = ElectricalCharge(system.GetSpecies(system.GetIndexOfPrimarySpecies(iPrimary)).GetName());
	
	ze.resize(numEquilibrium);
	
	for(unsigned iEquilibrium = 0; iEquilibrium < numEquilibrium; ++iEquilibrium)
		ze[iEquilibrium] = ElectricalCharge(system.GetSpecies(system.GetIndexOfEquilibriumSpecies(iEquilibrium)).GetName());
	
	Cj.resize(numPrimary);
	
	for(unsigned iPrimary = 0; iPrimary < numPrimary; ++iPrimary)
		Cj[iPrimary] = system.GetMolesOfSpecies(system.GetIndexOfPrimarySpecies(iPrimary));
		
	Ce.resize(numEquilibrium);
	
	for(unsigned iEquilibrium = 0; iEquilibrium < numEquilibrium; ++iEquilibrium)
		Ce[iEquilibrium] = system.GetMolesOfSpecies(system.GetIndexOfEquilibriumSpecies(iEquilibrium));
	
	system.UpdateSystem();
	
	aj.resize(numPrimary);
	
	for(unsigned iPrimary = 0; iPrimary < numPrimary; ++iPrimary)
		aj[iPrimary] = system.GetActivityOfSpecies(system.GetIndexOfPrimarySpecies(iPrimary));
	
	ae.resize(numEquilibrium);
	
	for(unsigned iEquilibrium = 0; iEquilibrium < numEquilibrium; ++iEquilibrium)
		ae[iEquilibrium] = system.GetActivityOfSpecies(system.GetIndexOfEquilibriumSpecies(iEquilibrium));
	
	Qe.resize(numEquilibrium);
	
	for(unsigned iEquilibrium = 0; iEquilibrium < numEquilibrium; ++iEquilibrium)
		Qe[iEquilibrium] = system.GetEquilibriumReactionQuotient(iEquilibrium);
	
	Ke.resize(numEquilibrium);
	
	const double T = system.GetTemperature();
	const double P = system.GetPressure();
	
	for(unsigned iEquilibrium = 0; iEquilibrium < numEquilibrium; ++iEquilibrium)
		Ke[iEquilibrium] = system.GetEquilibriumReaction(iEquilibrium).EquilibriumConstant(T, P);
	
	h.resize(numPrimary);
	
	alpha.resize(numPrimary, vector<double>(numPrimary));
	
	beta.resize(numPrimary, vector<double>(numEquilibrium));
	
	theta.resize(numPrimary, vector<double>(numPrimary));
	
	mu.resize(numPrimary, vector<double>(numEquilibrium));
	
	currentCondition = 0;
}

void Speciation::ImposeChargeBalance()
{
	if(currentCondition < system->GetNumOfPrimarySpecies())
	{
		for(unsigned iPrimary = 0; iPrimary < system->GetNumOfPrimarySpecies(); ++iPrimary)
			alpha[currentCondition][iPrimary] = zj[iPrimary];
	
		for(unsigned iEquilibrium = 0; iEquilibrium < system->GetNumOfEquilibriumSpecies(); ++iEquilibrium)
			beta[currentCondition][iEquilibrium] = ze[iEquilibrium];
		
		h[currentCondition] = 0.0;
		
		++currentCondition;
	}
	else
	{
		cout << "The condition for charge balance could not be imposed, since there are already sufficient conditions." << endl; // TODO: Create a log file to which these errors can be dumped in.
	}
}

void Speciation::ImposeUj(unsigned iPrimary, double value)
{
	if(currentCondition < system->GetNumOfPrimarySpecies())
	{
		h[currentCondition] = value;
		
		alpha[currentCondition][iPrimary] = 1.0;
		
		for(unsigned iEquilibrium = 0; iEquilibrium < system->GetNumOfEquilibriumSpecies(); ++iEquilibrium)
			beta[currentCondition][iEquilibrium] = system->GetEquilibriumReaction(iEquilibrium).GetStoichiometryOfSpecies(iPrimary);
		
		++currentCondition;
	}
	else
	{
		const Species& species = system->GetSpecies(system->GetIndexOfPrimarySpecies(iPrimary));
		
		cout << "The condition for Uj corresponding to species " << species.GetName() << " could not be imposed." << endl; // TODO: Create a log file to which these errors can be dumped in.
	}
}

void Speciation::ImposeCj(unsigned iPrimary, double value)
{
	if(currentCondition < system->GetNumOfPrimarySpecies())
	{
		h[currentCondition] = value;
		
		alpha[currentCondition][iPrimary] = 1.0;
		
		++currentCondition;
	}
	else
	{
		const Species& species = system->GetSpecies(system->GetIndexOfPrimarySpecies(iPrimary));
		
		cout << "The condition for Cj corresponding to species " << species.GetName() << " could not be imposed." << endl; // TODO: Create a log file to which these errors can be dumped in.
	}
}

void Speciation::ImposeCe(unsigned iEquilibrium, double value)
{
	if(currentCondition < system->GetNumOfPrimarySpecies())
	{
		h[currentCondition] = value;
		
		beta[currentCondition][iEquilibrium] = 1.0;
		
		++currentCondition;
	}
	else
	{
		const Species& species = system->GetSpecies(system->GetIndexOfPrimarySpecies(iEquilibrium));
		
		cout << "The condition for Ce corresponding to species " << species.GetName() << " could not be imposed." << endl; // TODO: Create a log file to which these errors can be dumped in.
	}
}

void Speciation::ImposeHj(unsigned iPrimary, double value)
{
	h[iPrimary] = value;
}

bool Speciation::Solve(double tolerance, unsigned maxiterations)
{
	// create a Newton-Raphson solver for the nonlinear system
	MultirootNewtonRaphson solver;
	
	// the number of primary species
	const unsigned numPrimary = system->GetNumOfPrimarySpecies();
	
	// the number of equilibrium species
	const unsigned numEquilibrium = system->GetNumOfEquilibriumSpecies();
	
	// the number of unknowns
	const unsigned numUnknowns = numPrimary + numEquilibrium;

	// set the number of equations in the non-linear system solver
	solver.SetNumEquations(numUnknowns);
	solver.SetTolerance(tolerance);

	vector<double> S(numUnknowns), oldS(numUnknowns);
	
	for(unsigned iPrimary = 0; iPrimary < numPrimary; ++iPrimary)
		S[iPrimary] = Cj[iPrimary];
	
	for(unsigned iEquilibrium = 0; iEquilibrium < numEquilibrium; ++iEquilibrium)
		S[iEquilibrium + numPrimary] = Ce[iEquilibrium];
	
	oldS = S;
	
	unsigned iter = 0;
	
//	cout << setw(10) << "Iter: "; // TODO be cleaned.
//	for(unsigned iSpecies = 0; iSpecies < system->GetNumOfSpecies(); ++iSpecies)
//		cout << setw(15) << system->GetSpecies(iSpecies).GetName();
//	cout << endl;
	
	do
	{
		for(unsigned i = 0; i < numUnknowns; ++i)
		{
			const double factor = max(1.0, - 2*(S[i] - oldS[i])/oldS[i]);
			
			S[i] = oldS[i] + 1.0/factor * (S[i] - oldS[i]);
		}
		
		oldS = S;
		
		UpdateComposition(S);
		
//		cout << setw(10) << iter; // TODO be cleaned.
//		for(unsigned iSpecies = 0; iSpecies < system->GetNumOfSpecies(); ++iSpecies)
//			cout << setw(15) << system->GetMolesOfSpecies(iSpecies);
//		cout << endl;
		
		// update the function data of the (nonlinear system) solver, since geosystem has been updated with new concentrations and activities for the aqueous species
		solver.SetFunction(bind(&Speciation::NonLinearSystemFunction, this, _1, _2));

		// update the jacobian data of the (nonlinear system) solver, since geosystem has been updated with new concentrations and activities for the aqueous species
		solver.SetJacobian(bind(&Speciation::NonLinearSystemJacobian, this, _1, _2));

		// update the function & jacobian data of the (nonlinear system) solver, since geosystem has been updated with new concentrations and activities for the aqueous species
		solver.SetFunctionJacobian(bind(&Speciation::NonLinearSystemFunctionJacobian, this, _1, _2, _3));
		
		if(++iter > maxiterations) return false;
	}
	while(solver.Iterate(S));
	
	UpdateComposition(S);
	
	return true;
}

void Speciation::NonLinearSystemFunction(const vector<double>& S, vector<double>& F)
{
	// the number of primary species in the system
	const unsigned numPrimary = system->GetNumOfPrimarySpecies();
	
	// the number of equilibrium species in the system
	const unsigned numEquilibrium = system->GetNumOfEquilibriumSpecies();
	
	for(unsigned iPrimary = 0; iPrimary < numPrimary; ++iPrimary)
	{
		F[iPrimary] = -h[iPrimary];
		
		F[iPrimary] += InnerProduct(alpha[iPrimary], Cj);
		
		F[iPrimary] += InnerProduct(beta[iPrimary],  Ce);
		
		F[iPrimary] += InnerProduct(theta[iPrimary], aj);
		
		F[iPrimary] += InnerProduct(mu[iPrimary],    ae);
	}
	
	for(unsigned iEquilibrium = 0; iEquilibrium < numEquilibrium; ++iEquilibrium)
	{
		//F[iEquilibrium + numPrimary] = Ke[iEquilibrium] - Qe[iEquilibrium];
		F[iEquilibrium + numPrimary] = log(Ke[iEquilibrium]/Qe[iEquilibrium]);
	}
}

void Speciation::NonLinearSystemJacobian(const vector<double>& S, vector<vector<double>>& J)
{
	// the number of primary species in the system
	const unsigned numPrimary = system->GetNumOfPrimarySpecies();
	
	// the number of equilibrium species in the system
	const unsigned numEquilibrium = system->GetNumOfEquilibriumSpecies();
	
	for(unsigned j = 0; j < numPrimary; ++j)
		for(unsigned i = 0; i < numPrimary; ++i)
			J[j][i] = alpha[j][i] + theta[j][i] * aj[i]/Cj[i];
	
	for(unsigned j = 0; j < numPrimary; ++j)
		for(unsigned e = 0; e < numEquilibrium; ++e)
			J[j][e + numPrimary] = beta[j][e] + mu[j][e] * ae[e]/Ce[e];
	
	for(unsigned e = 0; e < numEquilibrium; ++e)
	{
		for(unsigned iProduct = 0; iProduct < system->GetEquilibriumReaction(e).GetNumOfProducts(); ++iProduct)
		{
			const unsigned j = system->GetEquilibriumReaction(e).GetIndexOfProduct(iProduct);
			
			const double vej = system->GetEquilibriumReaction(e).GetStoichiometryOfProduct(iProduct);
			
			//J[e + numPrimary][j] = -vej * Qe[e]/Cj[j];
			J[e + numPrimary][j] = -vej/Cj[j];
		}
	}
	
	for(unsigned e = 0; e < numEquilibrium; ++e)
		//J[e + numPrimary][e + numPrimary] = Qe[e]/Ce[e];
		J[e + numPrimary][e + numPrimary] = 1.0/Ce[e];
}

void Speciation::NonLinearSystemFunctionJacobian(const vector<double>& S, vector<double>& F, vector<vector<double>>& J)
{
	NonLinearSystemFunction(S, F);
	NonLinearSystemJacobian(S, J);
}

void Speciation::UpdateComposition(const vector<double>& S)
{
	// the number of primary species
	const unsigned numPrimary = system->GetNumOfPrimarySpecies();
	
	// the number of equilibrium species
	const unsigned numEquilibrium = system->GetNumOfEquilibriumSpecies();
	
	for(unsigned iPrimary = 0; iPrimary < numPrimary; ++iPrimary)
		Cj[iPrimary] = S[iPrimary];
	
	for(unsigned iEquilibrium = 0; iEquilibrium < numEquilibrium; ++iEquilibrium)
		Ce[iEquilibrium] = S[iEquilibrium + numPrimary];
	
	system->SetMolesOfPrimarySpecies(Cj);
	
	system->SetMolesOfEquilibriumSpecies(Ce);
	
	system->UpdateSystem();
	
	for(unsigned iPrimary = 0; iPrimary < numPrimary; ++iPrimary)
		aj[iPrimary] = system->GetActivityOfSpecies(system->GetIndexOfPrimarySpecies(iPrimary));
	
	for(unsigned iEquilibrium = 0; iEquilibrium < numEquilibrium; ++iEquilibrium)
		ae[iEquilibrium] = system->GetActivityOfSpecies(system->GetIndexOfEquilibriumSpecies(iEquilibrium));
	
	for(unsigned iEquilibrium = 0; iEquilibrium < numEquilibrium; ++iEquilibrium)
		Qe[iEquilibrium] = system->GetEquilibriumReactionQuotient(iEquilibrium);
}
